from PIL import Image
import math
import os
import random


class VectorCompare(object):
    def magnitude(self, concordance):
        """
        计算向量的值
        :param concordance: 
        :return: 
        """
        total = 0
        for (word, count) in concordance.items():
            total += count ** 2
        return math.sqrt(total)

    def relation(self, concordance1, concordance2):
        """
        计算向量的cos值
        :param concordance1: 
        :param concordance2: 
        :return: 
        """
        top_value = 0
        for (word, count) in concordance1.items():
            if word in concordance2:
                top_value += count * concordance2[word]
        return top_value / (self.magnitude(concordance1) * self.magnitude(concordance2))


def build_vector(im):
    """
    生产向量
    :param im: 
    :return: 
    """
    d1 = {}
    count = 0
    for i in im.getdata():
        d1[count] = i
        count += 1

    return d1


v = VectorCompare()

icon_set = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '0', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j',
            'k',
            'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']

imageset = []

for letter in icon_set:
    for img in os.listdir('./toutiaoset/%s/' % (letter)):
        temp = []
        if img != "Thumbs.db" and img != ".DS_Store":  # windows check...
            temp.append(build_vector(Image.open("./toutiaoset/%s/%s" % (letter, img))))
        imageset.append({letter: temp})


def make_img(file_path):
    im = Image.open(file_path)
    im2 = Image.new("P", im.size, 255)
    im.convert("P")

    his = im.histogram()
    values = {}

    for i in range(256):
        values[i] = his[i]

    text_color = sorted(values.items(), key=lambda x: x[1], reverse=True)[:2][-1][0]

    temp = {}

    for x in range(im.size[1]):
        for y in range(im.size[0]):
            pix = im.getpixel((y, x))
            temp[pix] = pix
            if pix == text_color:  # these are the numbers to get
                im2.putpixel((y, x), 0)
    inletter = False
    foundletter = False
    start = 0
    end = 0

    letters = []

    for y in range(im2.size[0]):  # slice across
        for x in range(im2.size[1]):  # slice down
            pix = im2.getpixel((y, x))
            if pix != 255:
                inletter = True

        if foundletter is False and inletter is True:
            foundletter = True
            start = y

        if foundletter is True and inletter is False:
            foundletter = False
            end = y
            letters.append((start, end))

        inletter = False

    count = 0
    for letter in letters:
        im3 = im2.crop((letter[0], 0, letter[1], im2.size[1]))
        file_name = random.random() * 10
        # im3.save('./output/' + str(file_name) + '.gif')
        guess = []
        for image in imageset:
            for (x, y) in image.items():
                if len(y) != 0:
                    guess.append((v.relation(y[0], build_vector(im3)), x))

        guess.sort(reverse=True)
        print(guess[0])
        im3.save('./output/' + guess[0][1] + '___' + str(file_name) + '.gif')
        count += 1


for img in os.listdir('./cap/'):
    file = './cap/' + img
    make_img(file)
